import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

"""
将lora权重合并到大模型中
"""


def merge_lora_to_LLM():
    # model_name_or_path = "your_LLM_model_path"
    model_name_or_path = "D:/codes/nlp_about/pretrained_model/ycycyc02_chatglm3-6b"
    # adapter_name_or_path = "your_lora_model_path"
    adapter_name_or_path = "D:/codes/llm_about/self-llm/zzzzz_train/ChatGLM6B/output/ChatGLM/checkpoint-40"
    # save_path = "save_model_path"
    save_path = "D:/codes/llm_about/self-llm/zzzzz_train/ChatGLM6B/output/ChatGLM6BLoraD"

    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        trust_remote_code=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(model, adapter_name_or_path)
    model = model.merge_and_unload()

    tokenizer.save_pretrained(save_path)
    model.save_pretrained(save_path)


if __name__ == "__main__":
    merge_lora_to_LLM()
    print("done")
